#!/usr/bin/env python3
# -*- coding:utf-8 -*-
u"""
Created at 2020.01.23
A simple wrapper of ePSI pipeline

Thanks to Li Z.D. and Xia Q.Q.
"""
import logging
import multiprocessing as mp
import os
from glob import glob

import click
from tqdm import tqdm

import genomic
import tools

__dir__ = os.path.abspath(os.path.dirname(__file__))

logging.basicConfig(level=logging.DEBUG, format="%(asctime)s %(name)-12s %(levelname)-8s %(message)s")


def start_psi(args):
    tool = tools.StartPSI(*args)
    tool.start()


@click.command(
    context_settings=dict(help_option_names=['-h', '--help'])
)
@click.option(
    "-i",
    "--input-file",
    type=click.Path(exists=True),
    required=True,
    help="Path to input SJ.out.tab or directory contains SJ.out.tab"
)
@click.option(
    "-o",
    "--output",
    type=click.Path(),
    required=True,
    help="Path to input output directory or file"
)
@click.option(
    "-r",
    "--reference",
    type=click.Path(exists=True),
    required=True,
    help="Path to reference gtf file"
)
@click.option(
    "-p",
    "--processes",
    type=click.INT,
    default=1,
    help="CPU processes"
)
@click.option(
    "--read-length",
    type=click.INT,
    required=True,
    help="read length of bam file"
)
@click.option(
    "--prepare",
    is_flag=True,
    default=False,
    help="Indicates whether to prepare the gtf"
)
@click.option(
    "--bedtools",
    default="",
    help="customized bedtools path"
)
@click.option(
    "--bam-suffix",
    default="Aligned.sortedByCoord.out.bam",
    help="The suffix of STAR output bam file, script will use this to replace `SJ.out.tab` for locate bam files"
)
@click.version_option("0.1.0", message="Current version %(version)s")
def main(
        input_file,
        output,
        reference,
        processes,
        prepare,
        read_length,
        bedtools,
        bam_suffix
):
    u"""
    A python3 wrapper of ePSI pipeline

    This pipeline requires the bedtools v2.23 and htseq
    """

    input_file = os.path.abspath(input_file)
    output = os.path.abspath(output)
    reference = os.path.abspath(reference)
    bedtools = os.path.abspath(bedtools)

    # prepare output files
    output_directory, output_file = None, None
    if os.path.isfile(input_file):
        output_file = output
        output_directory = os.path.dirname(output_file)
        input_file = [input_file]
    else:
        output_directory = output
        input_file = glob(os.path.join(input_file, "*SJ.out.tab"))

    if output_directory and not os.path.exists(output_directory):
        os.makedirs(output_directory)

    if prepare:
        out_ref = os.path.join(output_directory, os.path.basename(reference))

        genomic.get_exonic_parts(reference, out_ref)

        reference = out_ref

    u"""
    gtf_file: str, 
    read_length: int, 
    junctions: str, 
    prefix: str,
    bedtools: str=""
    verbose=True
    """
    tasks = [
        [
            reference,  # exonic part gtf
            i.replace("SJ.out.tab", bam_suffix),  # bam file
            read_length,  # read length
            i,  # junctions
            os.path.join(output_directory, os.path.basename(i).replace("SJ.out.tab", "").strip(".")),  # output prefix
            bedtools  # path to bedtools
        ] for i in input_file
    ]

    mp.log_to_stderr(logging.DEBUG)
    with mp.Pool(processes) as p:
        list(tqdm(p.imap(start_psi, tasks), total=len(tasks)))


if __name__ == '__main__':
    main()
